"""
仿真信号生成演示
展示各种信号类型的生成和可视化
"""

import numpy as np
import matplotlib.pyplot as plt
from signal_generator import (
    SignalParams, HarmonicGenerator, ModulationGenerator,
    BearingFaultGenerator, CompositeSignalGenerator,
    plot_signal_analysis
)

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'Arial Unicode MS']
plt.rcParams['axes.unicode_minus'] = False


def demo_harmonic_signal():
    """演示谐波信号生成"""
    print("\n" + "="*60)
    print("1. 谐波信号生成演示")
    print("="*60)
    
    params = SignalParams(fs=10000, duration=1.0, noise_level=0.02)
    generator = HarmonicGenerator(params)
    
    # 生成包含5次谐波的信号，基频为50Hz
    signal, info = generator.generate(
        fundamental_freq=50.0,
        num_harmonics=5,
        amplitudes=[1.0, 0.5, 0.3, 0.2, 0.1]
    )
    
    print(f"\n信号类型: {info['type']}")
    print(f"基频: {info['fundamental_freq']} Hz")
    print(f"谐波频率: {info['frequencies']} Hz")
    print(f"谐波幅值: {info['amplitudes']}")
    
    # 绘图
    fig = plot_signal_analysis(generator.t, signal, params.fs, "谐波信号")
    plt.savefig('harmonic_signal.png', dpi=150, bbox_inches='tight')
    print("\n图像已保存: harmonic_signal.png")
    
    return signal, info


def demo_modulation_signal():
    """演示调制信号生成"""
    print("\n" + "="*60)
    print("2. 调制（边带）信号生成演示")
    print("="*60)
    
    params = SignalParams(fs=10000, duration=1.0, noise_level=0.01)
    generator = ModulationGenerator(params)
    
    # 生成调幅信号
    print("\n2.1 调幅(AM)信号")
    am_signal, am_info = generator.generate_am(
        carrier_freq=500.0,      # 载波频率 500Hz
        modulation_freq=30.0,    # 调制频率 30Hz
        carrier_amp=1.0,
        modulation_index=0.8
    )
    
    print(f"信号类型: {am_info['type']}")
    print(f"载波频率: {am_info['carrier_freq']} Hz")
    print(f"调制频率: {am_info['modulation_freq']} Hz")
    print(f"边带频率: {am_info['sideband_freqs']} Hz")
    
    fig = plot_signal_analysis(generator.t, am_signal, params.fs, "调幅信号 (AM)")
    plt.savefig('am_signal.png', dpi=150, bbox_inches='tight')
    print("图像已保存: am_signal.png")
    
    # 生成调频信号
    print("\n2.2 调频(FM)信号")
    fm_signal, fm_info = generator.generate_fm(
        carrier_freq=500.0,
        modulation_freq=30.0,
        carrier_amp=1.0,
        frequency_deviation=50.0
    )
    
    print(f"信号类型: {fm_info['type']}")
    print(f"载波频率: {fm_info['carrier_freq']} Hz")
    print(f"调制频率: {fm_info['modulation_freq']} Hz")
    print(f"频率偏移: {fm_info['frequency_deviation']} Hz")
    
    fig = plot_signal_analysis(generator.t, fm_signal, params.fs, "调频信号 (FM)")
    plt.savefig('fm_signal.png', dpi=150, bbox_inches='tight')
    print("图像已保存: fm_signal.png")
    
    return am_signal, fm_signal


def demo_bearing_fault_signal():
    """演示轴承故障信号生成"""
    print("\n" + "="*60)
    print("3. 轴承故障信号生成演示")
    print("="*60)
    
    params = SignalParams(fs=20000, duration=0.5, noise_level=0.05)
    generator = BearingFaultGenerator(params)
    
    rotation_freq = 25.0  # 转速 25Hz (1500 RPM)
    
    # 外圈故障
    print("\n3.1 外圈故障信号")
    outer_signal, outer_info = generator.generate_outer_race_fault(
        rotation_freq=rotation_freq,
        num_balls=9,
        impact_amplitude=2.0,
        damping=500.0,
        resonance_freq=3000.0
    )
    
    print(f"信号类型: {outer_info['type']}")
    print(f"转速: {outer_info['rotation_freq']} Hz ({outer_info['rotation_freq']*60:.0f} RPM)")
    print(f"外圈故障特征频率 (BPFO): {outer_info['BPFO']:.2f} Hz")
    print(f"共振频率: {outer_info['resonance_freq']} Hz")
    
    fig = plot_signal_analysis(generator.t, outer_signal, params.fs, "轴承外圈故障信号")
    plt.savefig('bearing_outer_fault.png', dpi=150, bbox_inches='tight')
    print("图像已保存: bearing_outer_fault.png")
    
    # 内圈故障
    print("\n3.2 内圈故障信号")
    inner_signal, inner_info = generator.generate_inner_race_fault(
        rotation_freq=rotation_freq,
        num_balls=9,
        impact_amplitude=2.0,
        damping=500.0,
        resonance_freq=3000.0
    )
    
    print(f"信号类型: {inner_info['type']}")
    print(f"转速: {inner_info['rotation_freq']} Hz ({inner_info['rotation_freq']*60:.0f} RPM)")
    print(f"内圈故障特征频率 (BPFI): {inner_info['BPFI']:.2f} Hz")
    print(f"共振频率: {inner_info['resonance_freq']} Hz")
    
    fig = plot_signal_analysis(generator.t, inner_signal, params.fs, "轴承内圈故障信号")
    plt.savefig('bearing_inner_fault.png', dpi=150, bbox_inches='tight')
    print("图像已保存: bearing_inner_fault.png")
    
    return outer_signal, inner_signal


def demo_composite_signal():
    """演示复合信号生成"""
    print("\n" + "="*60)
    print("4. 复合信号生成演示")
    print("="*60)
    print("\n复合信号包含:")
    print("  - 不平衡谐波 (转频及其倍频)")
    print("  - 齿轮啮合调制 (啮合频率被转频调制)")
    print("  - 轴承外圈故障 (周期性冲击)")
    
    params = SignalParams(fs=20000, duration=1.0, noise_level=0.08)
    generator = CompositeSignalGenerator(params)
    
    # 生成复合故障信号
    composite_signal, info = generator.generate_complex_fault_signal(
        rotation_freq=30.0  # 30Hz转速 (1800 RPM)
    )
    
    print(f"\n转速: 30 Hz (1800 RPM)")
    print(f"\n各组成部分:")
    
    harmonic_info = info['components']['harmonic']
    print(f"\n  谐波成分:")
    print(f"    频率: {harmonic_info['frequencies']} Hz")
    
    mod_info = info['components']['modulation']
    print(f"\n  调制成分 (齿轮啮合):")
    print(f"    载波频率: {mod_info['carrier_freq']} Hz")
    print(f"    调制频率: {mod_info['modulation_freq']} Hz")
    print(f"    边带频率: {mod_info['sideband_freqs']} Hz")
    
    bearing_info = info['components']['bearing_fault']
    print(f"\n  轴承故障成分:")
    print(f"    BPFO: {bearing_info['BPFO']:.2f} Hz")
    print(f"    共振频率: {bearing_info['resonance_freq']} Hz")
    
    # 绘图
    fig = plot_signal_analysis(generator.t, composite_signal, params.fs, "复合故障信号")
    plt.savefig('composite_signal.png', dpi=150, bbox_inches='tight')
    print("\n图像已保存: composite_signal.png")
    
    return composite_signal, info


def demo_signal_superposition():
    """演示自定义信号叠加"""
    print("\n" + "="*60)
    print("5. 自定义信号叠加演示")
    print("="*60)
    
    params = SignalParams(fs=10000, duration=1.0, noise_level=0.05)
    
    # 生成三个不同的信号
    print("\n生成三个独立信号:")
    
    # 信号1: 谐波
    gen1 = HarmonicGenerator(params)
    signal1, info1 = gen1.generate(fundamental_freq=60.0, num_harmonics=2)
    print(f"  信号1: 谐波信号 (60Hz)")
    
    # 信号2: 调制
    gen2 = ModulationGenerator(params)
    signal2, info2 = gen2.generate_am(carrier_freq=400.0, modulation_freq=25.0)
    print(f"  信号2: 调幅信号 (400Hz载波, 25Hz调制)")
    
    # 信号3: 轴承故障
    gen3 = BearingFaultGenerator(params)
    signal3, info3 = gen3.generate_outer_race_fault(rotation_freq=20.0)
    print(f"  信号3: 轴承外圈故障")
    
    # 叠加信号
    composite_gen = CompositeSignalGenerator(params)
    combined_signal = composite_gen.combine_signals(
        [signal1, signal2, signal3],
        weights=[1.0, 0.5, 1.5]  # 可以调整各信号的权重
    )
    
    print(f"\n信号已叠加，权重分别为: [1.0, 0.5, 1.5]")
    
    # 绘制对比图
    fig, axes = plt.subplots(5, 2, figsize=(14, 12))
    
    signals = [signal1, signal2, signal3, combined_signal]
    titles = ['谐波信号', '调幅信号', '轴承故障信号', '叠加信号']
    
    for i, (sig, title) in enumerate(zip(signals, titles)):
        # 时域
        t = params.duration * np.arange(len(sig)) / len(sig)
        axes[i, 0].plot(t[:2000], sig[:2000], linewidth=0.8)
        axes[i, 0].set_ylabel('幅值')
        axes[i, 0].set_title(f'{title} - 时域')
        axes[i, 0].grid(True, alpha=0.3)
        
        # 频域
        freq = np.fft.fftfreq(len(sig), 1/params.fs)
        fft_mag = np.abs(np.fft.fft(sig)) / len(sig)
        positive_idx = (freq > 0) & (freq < 2000)
        axes[i, 1].plot(freq[positive_idx], fft_mag[positive_idx], 'r-', linewidth=0.8)
        axes[i, 1].set_ylabel('幅值')
        axes[i, 1].set_title(f'{title} - 频域')
        axes[i, 1].grid(True, alpha=0.3)
    
    axes[-1, 0].set_xlabel('时间 (秒)')
    axes[-1, 1].set_xlabel('频率 (Hz)')
    
    # 添加说明
    explanation = "叠加效果: 可以看到最终信号包含了所有成分的特征"
    fig.text(0.5, 0.02, explanation, ha='center', fontsize=11, 
             bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))
    
    plt.tight_layout()
    plt.subplots_adjust(bottom=0.05)
    plt.savefig('signal_superposition.png', dpi=150, bbox_inches='tight')
    print("图像已保存: signal_superposition.png")
    
    return combined_signal


def main():
    """主函数"""
    print("="*60)
    print("仿真信号生成演示程序")
    print("基于MaintAGT论文的Sim2Real方法")
    print("="*60)
    
    # 运行各个演示
    demo_harmonic_signal()
    demo_modulation_signal()
    demo_bearing_fault_signal()
    demo_composite_signal()
    demo_signal_superposition()
    
    print("\n" + "="*60)
    print("所有演示完成！")
    print("="*60)
    print("\n生成的图像文件:")
    print("  1. harmonic_signal.png - 谐波信号")
    print("  2. am_signal.png - 调幅信号")
    print("  3. fm_signal.png - 调频信号")
    print("  4. bearing_outer_fault.png - 轴承外圈故障")
    print("  5. bearing_inner_fault.png - 轴承内圈故障")
    print("  6. composite_signal.png - 复合故障信号")
    print("  7. signal_superposition.png - 信号叠加对比")
    print("\n这些仿真信号可用于:")
    print("  - 故障诊断算法验证")
    print("  - 训练数据生成")
    print("  - 信号处理方法测试")
    print("  - 特征提取算法开发")
    

if __name__ == "__main__":
    main()


